import pygame
import random
import sys

# 初始化Pygame
pygame.init()

# 游戏常量
WIDTH, HEIGHT = 800, 600
GRID_SIZE = 20
GRID_WIDTH = WIDTH // GRID_SIZE
GRID_HEIGHT = HEIGHT // GRID_SIZE

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 255)

# 设置窗口
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("贪吃蛇游戏")

# 时钟控制游戏速度
clock = pygame.time.Clock()
FPS = 10

# 方向常量
UP = (0, -1)
DOWN = (0, 1)
LEFT = (-1, 0)
RIGHT = (1, 0)

class Snake:
    def __init__(self):
        self.reset()

    def reset(self):
        # 初始位置在屏幕中心
        self.position = (GRID_WIDTH // 2, GRID_HEIGHT // 2)
        # 初始长度为3
        self.segments = [self.position]
        # 初始方向向右
        self.direction = RIGHT
        # 初始分数
        self.score = 0
        # 是否游戏结束
        self.game_over = False

    def get_head_position(self):
        return self.segments[0]

    def turn(self, new_direction):
        # 防止直接反向移动
        if (new_direction[0] * -1, new_direction[1] * -1) == self.direction:
            return
        self.direction = new_direction

    def move(self):
        if self.game_over:
            return

        x, y = self.get_head_position()
        dx, dy = self.direction
        new_x = (x + dx) % GRID_WIDTH
        new_y = (y + dy) % GRID_HEIGHT
        new_position = (new_x, new_y)

        # 检测是否撞到自己
        if new_position in self.segments[1:]:
            self.game_over = True
            return

        self.segments.insert(0, new_position)
        # 如果吃到食物就增加长度，否则移除尾部
        if new_position == food.position:
            self.score += 10
            food.randomize_position()
        else:
            self.segments.pop()

    def draw(self):
        for segment in self.segments:
            x, y = segment
            rect = pygame.Rect((x * GRID_SIZE, y * GRID_SIZE), (GRID_SIZE - 1, GRID_SIZE - 1))
            pygame.draw.rect(screen, GREEN, rect)

        # 游戏结束时显示文字
        if self.game_over:
            font = pygame.font.SysFont(None, 72)
            text = font.render('游戏结束!', True, RED)
            screen.blit(text, (WIDTH//2 - text.get_width()//2, HEIGHT//2 - text.get_height()//2))
            font_small = pygame.font.SysFont(None, 36)
            text_small = font_small.render('按R键重新开始', True, WHITE)
            screen.blit(text_small, (WIDTH//2 - text_small.get_width()//2, HEIGHT//2 + text.get_height()))

class Food:
    def __init__(self):
        self.position = (0, 0)
        self.randomize_position()

    def randomize_position(self):
        # 确保食物不会出现在蛇身上
        while True:
            self.position = (random.randint(0, GRID_WIDTH - 1), random.randint(0, GRID_HEIGHT - 1))
            if self.position not in snake.segments:
                break

    def draw(self):
        x, y = self.position
        rect = pygame.Rect((x * GRID_SIZE, y * GRID_SIZE), (GRID_SIZE - 1, GRID_SIZE - 1))
        pygame.draw.rect(screen, RED, rect)

# 创建游戏对象
snake = Snake()
food = Food()

# 游戏主循环
running = True
while running:
    screen.fill(BLACK)

    # 事件处理
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP:
                snake.turn(UP)
            elif event.key == pygame.K_DOWN:
                snake.turn(DOWN)
            elif event.key == pygame.K_LEFT:
                snake.turn(LEFT)
            elif event.key == pygame.K_RIGHT:
                snake.turn(RIGHT)
            elif event.key == pygame.K_r and snake.game_over:
                snake.reset()
                food.randomize_position()
            elif event.key == pygame.K_ESCAPE:
                running = False

    # 游戏逻辑
    snake.move()

    # 绘制元素
    snake.draw()
    food.draw()

    # 显示分数
    font = pygame.font.SysFont(None, 36)
    score_text = font.render(f'分数: {snake.score}', True, WHITE)
    screen.blit(score_text, (10, 10))

    # 更新显示
    pygame.display.update()
    clock.tick(FPS)

# 退出游戏
pygame.quit()
sys.exit()